#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
@Author: yinwh
@Time: 2025-06-06 17:06
@File: VectorBase.py
@Version: 1.0.0
@Description: 
@Copyright: (c) 2025 by yinwh. All rights reserved.
"""

import os
from typing import List, Optional
import json
from Embeddings import BaseEmbeddings

import numpy as np
from tqdm import tqdm

class VectorStore:

    def __init__(self, document: List[str] = ['']) -> None:
        self.document = document

    def get_vector(self, EmbeddingModel: BaseEmbeddings) -> List[List[float]]:
        self.vector = []
        for doc in tqdm(self.document, desc="Calculating Embeddings"):
            self.vector.append(EmbeddingModel.get_embedding(doc))

        return self.vector

    def persist(self, path: str = "storage"):
        if not os.path.exists(path):
            os.makedirs(path)

        with open(f"{path}/document.json", "w", encoding="utf-8") as file:
            json.dump(self.document, file, ensure_ascii=False, indent=4)
        if self.vector:
            with open(f"{path}/vector.json", "w", encoding="utf-8") as file:
                json.dump(self.vector, file)

    def load_vector(self, path: str = "storage"):
        with open(f"{path}/vector.json", "r", encoding="utf-8") as f:
            self.vector = json.load(f)

        with open(f"{path}/document.json", "r", encoding="utf-8") as f:
            self.document = json.load(f)

    def get_similarity(self, vector1: List[float], vector2: List[float]) -> float:
        return BaseEmbeddings.cosine_similarity(vector1, vector2)

    def query(self, query: str, EmbeddingModel: BaseEmbeddings, k: int = 1) -> List[str]:
        # 获取查询向量
        query_vector = EmbeddingModel.get_embedding(query)
        # 计算余弦相似度，寻找相似的文档切片。

        similarity_scores = [self.get_similarity(query_vector, vector) for vector in self.vector]
        # 返回最相似的 k 个文档切片。
        return np.array(self.document)[np.argsort(similarity_scores)[-k:][::-1]].tolist()

    def sort_with_indices(self, arr: List[float]) -> tuple:
        """
        使用 argsort 对数组进行排序，并返回排序后的数组及其原始索引。

        :param arr: 需要排序的数组
        :return: 排序后的数组和原始索引
        """
        # 获取排序索引
        sorted_indices = np.argsort(arr)
        # 根据排序索引重新排列数组
        sorted_array = np.array(arr)[sorted_indices]
        # 返回排序后的数组和原始索引
        return sorted_array.tolist(), sorted_indices.tolist()

if __name__ == '__main__':
    vector_store = VectorStore()
    arr = [3.5, 1.2, 4.8, 2.1]
    sorted_array, original_indices = vector_store.sort_with_indices(arr)
    print("排序后的数组:", sorted_array)
    print("原始索引:", original_indices)
